{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DistributedTripletMarginLossMNIST.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "0SEqLaxOq7kw",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "eaca39d5-94ac-4e33-89a1-f97b30c66a27"
      },
      "source": [
        "!pip install pytorch-metric-learning\n",
        "!pip install faiss-cpu # we're using cpu for this example"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: pytorch-metric-learning in /usr/local/lib/python3.6/dist-packages (0.9.94)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning) (0.22.2.post1)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning) (1.7.0+cu101)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning) (0.8.1+cu101)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning) (1.18.5)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.6/dist-packages (from pytorch-metric-learning) (4.41.1)\n",
            "Requirement already satisfied: scipy>=0.17.0 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->pytorch-metric-learning) (1.4.1)\n",
            "Requirement already satisfied: joblib>=0.11 in /usr/local/lib/python3.6/dist-packages (from scikit-learn->pytorch-metric-learning) (0.17.0)\n",
            "Requirement already satisfied: dataclasses in /usr/local/lib/python3.6/dist-packages (from torch>=1.6.0->pytorch-metric-learning) (0.8)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.6/dist-packages (from torch>=1.6.0->pytorch-metric-learning) (3.7.4.3)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch>=1.6.0->pytorch-metric-learning) (0.16.0)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision->pytorch-metric-learning) (7.0.0)\n",
            "Requirement already satisfied: faiss-cpu in /usr/local/lib/python3.6/dist-packages (1.6.5)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vYmfxuw1tYdm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b414f826-84e6-4eb9-b80d-e9f406dfded1"
      },
      "source": [
        "######################################################################################################################\n",
        "### This script is modified from the guide on pytorch distributed training https://github.com/seba-1511/dist_tuto.pth/\n",
        "### https://pytorch.org/tutorials/intermediate/dist_tuto.html\n",
        "######################################################################################################################\n",
        "import os\n",
        "import torch\n",
        "import torch.distributed as dist\n",
        "from torch.nn.parallel import DistributedDataParallel as DDP\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "\n",
        "from math import ceil\n",
        "from random import Random\n",
        "from torch.multiprocessing import Process\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "from pytorch_metric_learning.utils import distributed as pml_dist\n",
        "from pytorch_metric_learning.utils.accuracy_calculator import AccuracyCalculator\n",
        "from pytorch_metric_learning import losses, miners, testers\n",
        "import numpy as np\n",
        "import logging\n",
        "logging.getLogger().setLevel(logging.INFO)\n",
        "\n",
        "\n",
        "class Partition(object):\n",
        "    \"\"\" Dataset-like object, but only access a subset of it. \"\"\"\n",
        "\n",
        "    def __init__(self, data, index):\n",
        "        self.data = data\n",
        "        self.index = index\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.index)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        data_idx = self.index[index]\n",
        "        return self.data[data_idx]\n",
        "\n",
        "\n",
        "class DataPartitioner(object):\n",
        "    \"\"\" Partitions a dataset into different chuncks. \"\"\"\n",
        "\n",
        "    def __init__(self, data, sizes=[0.7, 0.2, 0.1], seed=1234):\n",
        "        self.data = data\n",
        "        self.partitions = []\n",
        "        rng = Random()\n",
        "        rng.seed(seed)\n",
        "        data_len = len(data)\n",
        "        indexes = [x for x in range(0, data_len)]\n",
        "        rng.shuffle(indexes)\n",
        "\n",
        "        for frac in sizes:\n",
        "            part_len = int(frac * data_len)\n",
        "            self.partitions.append(indexes[0:part_len])\n",
        "            indexes = indexes[part_len:]\n",
        "\n",
        "    def use(self, partition):\n",
        "        return Partition(self.data, self.partitions[partition])\n",
        "\n",
        "\n",
        "class Net(nn.Module):\n",
        "    \"\"\" Network architecture. \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(Net, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)\n",
        "        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)\n",
        "        self.conv2_drop = nn.Dropout2d()\n",
        "        self.fc1 = nn.Linear(320, 50)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(F.max_pool2d(self.conv1(x), 2))\n",
        "        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))\n",
        "        x = x.view(-1, 320)\n",
        "        return self.fc1(x)\n",
        "\n",
        "def get_MNIST(train):\n",
        "    return datasets.MNIST(\n",
        "            './data',\n",
        "            train=train,\n",
        "            download=True,\n",
        "            transform=transforms.Compose([\n",
        "                transforms.ToTensor(),\n",
        "                transforms.Normalize((0.1307, ), (0.3081, ))\n",
        "            ]))\n",
        "    \n",
        "\n",
        "def partition_dataset(dataset):\n",
        "    \"\"\" Partitioning MNIST \"\"\"\n",
        "    size = dist.get_world_size()\n",
        "    bsz = 512 // size\n",
        "    partition_sizes = [1.0 / size for _ in range(size)]\n",
        "    partition = DataPartitioner(dataset, partition_sizes)\n",
        "    partition = partition.use(dist.get_rank())\n",
        "    train_set = torch.utils.data.DataLoader(\n",
        "        partition, batch_size=bsz, shuffle=True)\n",
        "    return train_set, bsz\n",
        "\n",
        "\n",
        "### convenient function from pytorch-metric-learning ###\n",
        "def get_all_embeddings(dataset, model, data_device):\n",
        "    # dataloader_num_workers has to be 0 to avoid pid error\n",
        "    # This only happens when within multiprocessing\n",
        "    tester = testers.BaseTester(dataloader_num_workers=0, data_device=data_device)\n",
        "    return tester.get_all_embeddings(dataset, model)\n",
        "\n",
        "\n",
        "### compute accuracy using AccuracyCalculator from pytorch-metric-learning ###\n",
        "def test(train_set, test_set, model, accuracy_calculator, data_device):\n",
        "    train_embeddings, train_labels = get_all_embeddings(train_set, model, data_device)\n",
        "    test_embeddings, test_labels = get_all_embeddings(test_set, model, data_device)\n",
        "    print(\"Computing accuracy\")\n",
        "    accuracies = accuracy_calculator.get_accuracy(test_embeddings, \n",
        "                                                train_embeddings,\n",
        "                                                test_labels,\n",
        "                                                train_labels,\n",
        "                                                False)\n",
        "    print(\"Validation set accuracy (Precision@1) = {}\".format(accuracies[\"precision_at_1\"]))\n",
        "\n",
        "def test_model(rank, train_set, test_set, model, epoch, data_device):\n",
        "    if rank == 0:\n",
        "        print(\"Computing validation set accuracy for epoch {}\".format(epoch))\n",
        "        accuracy_calculator = AccuracyCalculator(include = (\"precision_at_1\",), k = 1)\n",
        "        test(train_set, test_set, model, accuracy_calculator, data_device)\n",
        "    dist.barrier()\n",
        "\n",
        "\n",
        "def run(rank, size, train_dataset, val_dataset):\n",
        "    \"\"\" Distributed Synchronous SGD Example \"\"\"\n",
        "    print(\"Rank {} entering the 'run' function\".format(rank))\n",
        "    torch.manual_seed(1234)\n",
        "    train_set, bsz = partition_dataset(train_dataset)\n",
        "    dist.barrier()\n",
        "    ### use this if you have multiple GPUs ###\n",
        "    # device = torch.device(\"cuda:{}\".format(rank))\n",
        "    device = torch.device(\"cpu\")\n",
        "    model = Net()\n",
        "    ### if you have multiple GPUs, set this to DDP(model.to(device), device_ids=[rank])\n",
        "    model = DDP(model.to(device))\n",
        "    test_model(rank, train_dataset, val_dataset, model, \"untrained\", device)\n",
        "\n",
        "    optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
        "\n",
        "    #####################################\n",
        "    ### pytorch-metric-learning stuff ###\n",
        "    loss_fn = losses.TripletMarginLoss()\n",
        "    ### if you have multiple GPUs, set this to pml_dist.DistributedLossWrapper(loss=loss_fn, device_ids=[rank])\n",
        "    loss_fn = pml_dist.DistributedLossWrapper(loss=loss_fn)\n",
        "    miner = miners.MultiSimilarityMiner()\n",
        "    miner = pml_dist.DistributedMinerWrapper(miner=miner)\n",
        "    ### pytorch-metric-learning stuff ###\n",
        "    #####################################\n",
        "\n",
        "    num_batches = ceil(len(train_set.dataset) / float(bsz))\n",
        "    for epoch in range(1):\n",
        "        epoch_loss = 0.0\n",
        "        print(\"Rank {} starting epoch {}\".format(rank, epoch))\n",
        "        for i, (data, target) in enumerate(train_set):\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            optimizer.zero_grad()\n",
        "            output = model(data)\n",
        "            hard_pairs = miner(output, target)\n",
        "            loss = loss_fn(output, target, hard_pairs)\n",
        "            epoch_loss += loss.item()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            if i % 10 == 0:\n",
        "                print('Rank {}, iteration {}, loss {}, num pos pairs {}, num neg pairs {}'.\\\n",
        "                      format(rank, i, loss.item(), miner.miner.num_pos_pairs, miner.miner.num_neg_pairs))\n",
        "            dist.barrier()\n",
        "\n",
        "        print('Rank {}, epoch {}, average loss {}'.format(rank, epoch, epoch_loss/num_batches))\n",
        "        test_model(rank, train_dataset, val_dataset, model, epoch, device)\n",
        "\n",
        "\n",
        "\n",
        "#######################################\n",
        "### Set backend='nccl' if using GPU ###\n",
        "#######################################\n",
        "def init_processes(rank, size, fn, train_dataset, val_dataset, backend='gloo'):\n",
        "    \"\"\" Initialize the distributed environment. \"\"\"\n",
        "    os.environ['MASTER_ADDR'] = 'localhost'\n",
        "    os.environ['MASTER_PORT'] = '29500'\n",
        "    dist.init_process_group(backend, rank=rank, world_size=size)\n",
        "    fn(rank, size, train_dataset, val_dataset)\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    train_dataset = get_MNIST(True)\n",
        "    val_dataset = get_MNIST(False)\n",
        "\n",
        "    size = 4\n",
        "    processes = []\n",
        "    for rank in range(size):\n",
        "        p = Process(target=init_processes, args=(rank, size, run, train_dataset, val_dataset))\n",
        "        p.start()\n",
        "        processes.append(p)\n",
        "\n",
        "    for p in processes:\n",
        "        p.join()"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Rank 0 entering the 'run' function\n",
            "Rank 2 entering the 'run' function\n",
            "Rank 3 entering the 'run' function\n",
            "Rank 1 entering the 'run' function\n",
            "Computing validation set accuracy for epoch untrained\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1875/1875 [00:22<00:00, 83.58it/s]\n",
            "100%|██████████| 313/313 [00:03<00:00, 84.06it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Computing accuracy\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:root:running k-nn with k=1\n",
            "INFO:root:embedding dimensionality is 50\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Validation set accuracy (Precision@1) = 0.9228\n",
            "Rank 0 starting epoch 0\n",
            "Rank 1 starting epoch 0\n",
            "Rank 3 starting epoch 0\n",
            "Rank 2 starting epoch 0\n",
            "Rank 2, iteration 0, loss 0.22071301937103271, num pos pairs 26684, num neg pairs 233135\n",
            "Rank 3, iteration 0, loss 0.22071301937103271, num pos pairs 26684, num neg pairs 233135\n",
            "Rank 1, iteration 0, loss 0.22071301937103271, num pos pairs 26684, num neg pairs 233135\n",
            "Rank 0, iteration 0, loss 0.22071301937103271, num pos pairs 26684, num neg pairs 233135\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n",
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n",
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n",
            "INFO:root:Reducer buckets have been rebuilt in this iteration.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Rank 0, iteration 10, loss 0.052407171577215195, num pos pairs 25802, num neg pairs 235806\n",
            "Rank 1, iteration 10, loss 0.052407171577215195, num pos pairs 25802, num neg pairs 235806\n",
            "Rank 3, iteration 10, loss 0.052407171577215195, num pos pairs 25802, num neg pairs 235806\n",
            "Rank 2, iteration 10, loss 0.052407171577215195, num pos pairs 25802, num neg pairs 235806\n",
            "Rank 1, iteration 20, loss 0.05540177971124649, num pos pairs 26076, num neg pairs 234137\n",
            "Rank 2, iteration 20, loss 0.05540177971124649, num pos pairs 26076, num neg pairs 234137\n",
            "Rank 3, iteration 20, loss 0.05540177971124649, num pos pairs 26076, num neg pairs 234137\n",
            "Rank 0, iteration 20, loss 0.05540177971124649, num pos pairs 26076, num neg pairs 234137\n",
            "Rank 1, iteration 30, loss 0.055692050606012344, num pos pairs 26650, num neg pairs 233889\n",
            "Rank 3, iteration 30, loss 0.055692050606012344, num pos pairs 26650, num neg pairs 233889\n",
            "Rank 2, iteration 30, loss 0.055692050606012344, num pos pairs 26650, num neg pairs 233889\n",
            "Rank 0, iteration 30, loss 0.055692050606012344, num pos pairs 26650, num neg pairs 233889\n",
            "Rank 3, iteration 40, loss 0.0504874549806118, num pos pairs 26008, num neg pairs 233677\n",
            "Rank 0, iteration 40, loss 0.0504874549806118, num pos pairs 26008, num neg pairs 233677\n",
            "Rank 1, iteration 40, loss 0.0504874549806118, num pos pairs 26008, num neg pairs 233677\n",
            "Rank 2, iteration 40, loss 0.0504874549806118, num pos pairs 26008, num neg pairs 233677\n",
            "Rank 1, iteration 50, loss 0.05479405075311661, num pos pairs 26312, num neg pairs 231127\n",
            "Rank 2, iteration 50, loss 0.05479405075311661, num pos pairs 26312, num neg pairs 231127\n",
            "Rank 3, iteration 50, loss 0.05479405075311661, num pos pairs 26312, num neg pairs 231127\n",
            "Rank 0, iteration 50, loss 0.05479405075311661, num pos pairs 26312, num neg pairs 231127\n",
            "Rank 2, iteration 60, loss 0.055632807314395905, num pos pairs 26042, num neg pairs 227909\n",
            "Rank 1, iteration 60, loss 0.055632807314395905, num pos pairs 26042, num neg pairs 227909\n",
            "Rank 0, iteration 60, loss 0.055632807314395905, num pos pairs 26042, num neg pairs 227909\n",
            "Rank 3, iteration 60, loss 0.055632807314395905, num pos pairs 26042, num neg pairs 227909\n",
            "Rank 3, iteration 70, loss 0.05833852291107178, num pos pairs 26428, num neg pairs 226380\n",
            "Rank 1, iteration 70, loss 0.05833852291107178, num pos pairs 26428, num neg pairs 226380\n",
            "Rank 2, iteration 70, loss 0.05833852291107178, num pos pairs 26428, num neg pairs 226380\n",
            "Rank 0, iteration 70, loss 0.05833852291107178, num pos pairs 26428, num neg pairs 226380\n",
            "Rank 0, iteration 80, loss 0.05444473400712013, num pos pairs 26222, num neg pairs 228766\n",
            "Rank 1, iteration 80, loss 0.05444473400712013, num pos pairs 26222, num neg pairs 228766\n",
            "Rank 3, iteration 80, loss 0.05444473400712013, num pos pairs 26222, num neg pairs 228766\n",
            "Rank 2, iteration 80, loss 0.05444473400712013, num pos pairs 26222, num neg pairs 228766\n",
            "Rank 3, iteration 90, loss 0.05678363889455795, num pos pairs 25966, num neg pairs 219641\n",
            "Rank 1, iteration 90, loss 0.05678363889455795, num pos pairs 25966, num neg pairs 219641\n",
            "Rank 2, iteration 90, loss 0.05678363889455795, num pos pairs 25966, num neg pairs 219641\n",
            "Rank 0, iteration 90, loss 0.05678363889455795, num pos pairs 25966, num neg pairs 219641\n",
            "Rank 3, iteration 100, loss 0.05184846743941307, num pos pairs 25930, num neg pairs 220450\n",
            "Rank 1, iteration 100, loss 0.05184846743941307, num pos pairs 25930, num neg pairs 220450\n",
            "Rank 0, iteration 100, loss 0.05184846743941307, num pos pairs 25930, num neg pairs 220450\n",
            "Rank 2, iteration 100, loss 0.05184846743941307, num pos pairs 25930, num neg pairs 220450\n",
            "Rank 0, iteration 110, loss 0.06211327761411667, num pos pairs 25980, num neg pairs 213528\n",
            "Rank 3, iteration 110, loss 0.06211327761411667, num pos pairs 25980, num neg pairs 213528\n",
            "Rank 1, iteration 110, loss 0.06211327761411667, num pos pairs 25980, num neg pairs 213528\n",
            "Rank 2, iteration 110, loss 0.06211327761411667, num pos pairs 25980, num neg pairs 213528\n",
            "Rank 0, epoch 0, average loss 0.05798070022236493\n",
            "Rank 2, epoch 0, average loss 0.05798070022236493\n",
            "Rank 3, epoch 0, average loss 0.05798070022236493\n",
            "Rank 1, epoch 0, average loss 0.05798070022236493\n",
            "Computing validation set accuracy for epoch 0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1875/1875 [00:22<00:00, 83.58it/s]\n",
            "100%|██████████| 313/313 [00:03<00:00, 80.88it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Computing accuracy\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "INFO:root:running k-nn with k=1\n",
            "INFO:root:embedding dimensionality is 50\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Validation set accuracy (Precision@1) = 0.9763\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}